module tb_key_expander();

wire [127:0] key_out;
reg  [127:0] key_in, key_golden;

wire [31:0] sbox_in, sbox_out;

reg [4:0] round;
reg enc_dec;


key_expander KEY_EXPANDER
(
	.key_out ( key_out ),
	.g_in		 ( sbox_in ),
	.g_out   ( sbox_out),
	.key_in  ( key_in  ),
	.round   ( round[3:0]),
	.enc_dec ( enc_dec)
);

sBox SBOX32
  (
    .sbox_out_enc ( sbox_out ),
		.sbox_out_dec (),
    .sbox_in      ( sbox_in  ),
    .enc_dec      ( 1'b1  )
  );


initial
	begin
		enc_dec = 1;
		//key_in = {32'h3c_4f_cf_09, 32'h88_15_f7_ab, 32'ha6_d2_ae_28, 32'h16_15_7e_2b};
		key_golden = {32'h2b_7e_15_16, 32'h28_ae_d2_a6, 32'hab_f7_15_88, 32'h09_cf_4f_3c};
		key_in = key_golden;
		for(round = 0; round < 10; round = round + 1)
			begin
				#5;
				key_in[32*(3 + 1) - 1 : 32*3] = key_out[32*(3 + 1) - 1 : 32*3];
				#5;
				key_in[32*(2 + 1) - 1 : 32*2] = key_out[32*(2 + 1) - 1 : 32*2];
				#5;
				key_in[32*(1 + 1) - 1 : 32*1] = key_out[32*(1 + 1) - 1 : 32*1];
				#5;
				key_in[32*(0 + 1) - 1 : 32*0] = key_out[32*(0 + 1) - 1 : 32*0];
			end
		enc_dec = 0;
			for(round = 9; round > 0; round = round - 1)
				begin
					#5
					key_in = key_out;
				end
				#5
				if(key_out != key_golden)
					begin
						$display("TEST FAILL!!");
						$stop;
					end
				$display("TEST PASSED!!!");
				$stop;
	end

endmodule
